package leetcodecompetition.D373;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

/**
 * @Author: Siantar
 * @Date: 2023-11-26-10:40
 * @Description: 1.0
 */
public class T100132 {
    private static final int AEIOU_MASK = 1065233;

    public long beautifulSubstrings(String s, int k) {
        k = pSqrt(k * 4);
        Map<Integer, Integer> cnt = new HashMap<>();
        int n = s.length();
        int sum = n; // 保证非负
        cnt.put((k - 1) << 16 | sum, 1); // 添加 (k-1, sum)
        long ans = 0;
        for (int i = 0; i < n; i++) {
            int bit = (AEIOU_MASK >> (s.charAt(i) - 'a')) & 1;
            sum += bit * 2 - 1; // 1 -> 1    0 -> -1
            ans += cnt.merge((i % k) << 16 | sum, 1, Integer::sum) - 1; // ans += cnt[(i%k,sum)]++
        }
        return ans;
    }

    private int pSqrt(int n) {
        int res = 1;
        for (int i = 2; i * i <= n; i++) {
            int i2 = i * i;
            while (n % i2 == 0) {
                res *= i;
                n /= i2;
            }
            if (n % i == 0) {
                res *= i;
                n /= i;
            }
        }
        if (n > 1) {
            res *= n;
        }
        return res;
    }
}
